#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>
#include <mpi.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_sort_float.h>
#include <gsl/gsl_heapsort.h>
#ifdef USEOPENMP
#include <omp.h>
#endif
#include <sys/time.h>

#include "halofinder.h"

static double wtime(void)
{
  struct timeval tp;
  gettimeofday(&tp,NULL);
  return ((double) (tp.tv_sec)) + ((double) (tp.tv_usec))/1e6;
}

void computehaloradii(void)
{
  int i,j,periodic,NumNbrs,maxNumNbrs,maxVcircInd;
  float searchRadius;
  long NumSindex;
  size_t *sindex,*tmpsindex;
  float haloMass,haloRadius,maxVCirc;
  float *nbrsRad2;
  int *nbrsInd;
  int foundHaloMass;
  int NumToDo,NumDone;
  float vel[3];
  double time,diff;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"computing halo radii\n");
#endif
  
#ifndef LIGHTCONE
  periodic = 1;
#else
  periodic = 0;
#endif
  
#ifdef USEOPENMP
#pragma omp parallel default(none) shared(ThisTask,stderr,NumHalos,Halos,PartsPx,PartsPy,PartsPz,PartsVx,PartsVy,PartsVz,PartskdTreeData,haloFinderData,periodic) \
  private(i,sindex,NumSindex,maxNumNbrs,nbrsRad2,nbrsInd,NumNbrs,searchRadius,haloMass,tmpsindex,maxVCirc,maxVcircInd,foundHaloMass,j,haloRadius,NumToDo,NumDone,vel,time,diff)
#endif
  {
    //setup
    sindex = NULL;
    NumSindex = 0;
    
    //set max num nbrs
    maxNumNbrs = 100000;
    nbrsRad2 = (float*)malloc(sizeof(float)*maxNumNbrs);
    assert(nbrsRad2 != NULL);
    nbrsInd = (int*)malloc(sizeof(int)*maxNumNbrs);
    assert(nbrsInd != NULL);
    
#ifdef USEOPENMP
    NumToDo = NumHalos/omp_get_num_threads();
#else
    NumToDo = NumHalos;
#endif
    NumDone = 0;
    
    time = wtime();
    
#ifdef USEOPENMP
#pragma omp for schedule(guided)
#endif
    for(i=0;i<NumHalos;++i)
      {
#ifdef DEBUG
	if(NumDone%(NumToDo/100) == 0 && ThisTask == 0)
	  {
	    diff = wtime()-time;
	    fprintf(stderr,"%d: %d of %d (%0.2lf percent) done in computing halo radii, %lf halos per second\n",
		    ThisTask,NumDone,NumToDo,((float) NumDone)/((float) NumToDo)*100.0,
		    ((double) (NumDone))/diff);
	  }
#endif
	++NumDone;
	
	//get radius of meanOverdensRefineHaloCenter for Halo
	searchRadius = 0.3;
	NumNbrs = get_nnbrs_kdtree(Halos[i].pos,searchRadius,periodic,haloFinderData.BoxLengthSnapshot,
				   &nbrsRad2,&nbrsInd,&maxNumNbrs,PartsPx,PartsPy,PartsPz,PartskdTreeData);
	haloMass = (float) (NumNbrs);
	while(searchRadius <= halo_truncation_radius(haloMass*haloFinderData.PartMass))
	  {
	    searchRadius = (float) (searchRadius*1.25);
	    NumNbrs = get_nnbrs_kdtree(Halos[i].pos,searchRadius,periodic,haloFinderData.BoxLengthSnapshot,
				       &nbrsRad2,&nbrsInd,&maxNumNbrs,PartsPx,PartsPy,PartsPz,PartskdTreeData);
	    haloMass = (float) (NumNbrs);
	  }
	if(NumNbrs > NumSindex)
	  {
	    tmpsindex = (size_t*)realloc(sindex,sizeof(size_t)*(NumNbrs));
	    assert(tmpsindex != NULL);
	    sindex = tmpsindex;
	    NumSindex = NumNbrs;
	  }
	gsl_sort_float_index(sindex,nbrsRad2,(size_t) 1,(size_t) NumNbrs);
	
	//work from outside in until radius and mass match
	maxVCirc = (float) (sqrt(BIG_G*haloFinderData.PartMass/sqrt(nbrsRad2[sindex[NumNbrs-1]])));
	maxVcircInd = NumNbrs-1;
	foundHaloMass = 0;
	for(j=NumNbrs-1;j>=0;--j)
	  {
	    haloRadius = (float) sqrt(nbrsRad2[sindex[j]]);
	    haloMass = (float) (j+1);
	    
	    if(sqrt(BIG_G*haloMass*haloFinderData.PartMass/haloRadius) > maxVCirc && haloRadius > 3.0*haloFinderData.SofteningLength)
	      {
		maxVCirc = (float) (sqrt(BIG_G*haloMass*haloFinderData.PartMass/haloRadius));
		maxVcircInd = j;
	      }
	    
	    if(haloRadius <= halo_truncation_radius(haloMass*haloFinderData.PartMass) && foundHaloMass == 0)
	      {
		foundHaloMass = 1;
		Halos[i].mass = haloMass;
	      }
	  }
	Halos[i].vcirc = maxVCirc;
	Halos[i].radius = halo_truncation_radius(Halos[i].mass*haloFinderData.PartMass);
	
	//compute center of mass velocity
	NumNbrs = maxVcircInd+1;
	vel[0] = 0.0;
	vel[1] = 0.0;
	vel[2] = 0.0;
	for(j=0;j<NumNbrs;++j)
	  {
	    vel[0] += PartsVx[nbrsInd[sindex[j]]];
	    vel[1] += PartsVy[nbrsInd[sindex[j]]];
	    vel[2] += PartsVz[nbrsInd[sindex[j]]];
	  }
	vel[0] /= ((float) NumNbrs);
	vel[1] /= ((float) NumNbrs);
	vel[2] /= ((float) NumNbrs);
	
	Halos[i].vel[0] = vel[0];
	Halos[i].vel[1] = vel[1];
	Halos[i].vel[2] = vel[2];
      }
    
    if(NumSindex > 0)
      free(sindex);
    free(nbrsRad2);
    free(nbrsInd);
  } /* end of parallel region */
  
#ifdef DEBUG
#if DEBUG_LEVEL > 0
  if(ThisTask == 0)
    fprintf(stderr,"done computing halo radii\n");
#endif
#endif
}

